
#include "sinh_custom_tiling.h"
#include "register/op_def_registry.h"


namespace optiling
{
    constexpr uint32_t BLOCK_DIM = 24;                        
    constexpr uint32_t SIZE_OF_HALF = 2;                      
    constexpr uint32_t BLOCK_SIZE = 32;                       
    constexpr uint32_t ALIGN_NUM = BLOCK_SIZE / SIZE_OF_HALF; 
    static ge::graphStatus TilingFunc(gert::TilingContext *context)
    {

        TilingData tiling;
        uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();
        context->SetBlockDim(BLOCK_DIM);

        
        uint32_t totalLengthAligned = ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
        
        uint32_t formerNum = (totalLengthAligned / ALIGN_NUM) % BLOCK_DIM;
        uint32_t tailNum = BLOCK_DIM - formerNum;
        
        uint32_t formerLength = ((totalLengthAligned / BLOCK_DIM + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
        uint32_t tailLength = (totalLengthAligned / BLOCK_DIM / ALIGN_NUM) * ALIGN_NUM;

        
        tiling.set_formerNum(formerNum);
        tiling.set_tailNum(tailNum);
        tiling.set_formerLength(formerLength);
        tiling.set_tailLength(tailLength);
        tiling.set_alignNum(ALIGN_NUM);

        
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
        context->SetTilingKey(1);
        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;

        return ge::GRAPH_SUCCESS;
    }
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class SinhCustom : public OpDef {
public:
    explicit SinhCustom(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910");

    }
};

OP_ADD(SinhCustom);
}
